// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file 
\brief Implementation of Sample Array/Individual-Major Allelic Data 
*/

#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include "copyfile.h"
#include "individual.h"
#include "locus.h"
#include "pedfile.h"
#include "pmaths.h"
#include "suffixes.h"
#include "utils.h"

#define msg_BUFFER_LENGTH 1024
#define DATA_BUFFER_SIZE 80
#define INITIAL_ITEM_ARRAY_SIZE 1000

int pedfile_read(struct pedfile *ped, struct mapfile *map, struct selected_options *ops) 
{
	FILE* fin = NULL;
	BOOL isPedOpen = FALSE, doLoop1 = FALSE, inComment = FALSE;
	char msg[msg_BUFFER_LENGTH];
	struct individual **ambiguous = NULL;
	int nAmbiguous = 0, ambiguous_available = 0;
	int nmale = 0, nfemale = 0;
	char buffer[DATA_BUFFER_SIZE];
	int c = 0; // number of individuals
	int nm = 0; // No. none-missing phenotypes.
	int j;
	char _chr;
	
	if (ped == NULL || map == NULL || ops == NULL) return 0;
	if ((fin = fopen(ops->szPedFileName, "r")) != NULL)  isPedOpen = TRUE;
	if (!isPedOpen) {
		sprintf(msg, "Problem opening PED file, errno = %i",errno);
		if (ops->verbose) error(msg);
		return 0;
	}

	// Stream through the file once to get the number of people in the sample
	// data.
	_chr = (char) fgetc(fin);
	while (_chr != EOF) {
		_chr = (char) fgetc(fin);
		if (_chr == '#') {
			inComment = TRUE;
		} else if (_chr == '\n') {
			if (inComment) inComment = FALSE;
			else ped->nSample++;
		}
	}
	if (fin != NULL) fclose(fin);

	if (ped->nSample == 0) {
		sprintf(msg, "Inputted PED file [%s] contains on records.", ops->szPedFileName);
		if (ops->verbose) error(msg);
		return 0;
	}

	if (!pedfile_alloc_sample(ped, ped->nSample)) {
		sprintf(msg, "Unable to allocate Inputted PED file [%s] storage array.", ops->szPedFileName);
		if (ops->verbose) error(msg);
		return 0;
	}
	if (ops->isSnpMajor) {
		int i;
		for (i = 0; i < map->nSNP; i++) {
			if (!csnp_alloc_one(map->SNP[i], ped->nSample) || 
				!csnp_alloc_two(map->SNP[i], ped->nSample)) 
			{
				sprintf(msg, "When in SNP major mode, unable to assign the results arrays", ops->szPedFileName);
				if (ops->verbose) error(msg);
				return 0;
			}
		}
	}

	// Stream in the PED file.
	fin = fopen(ops->szPedFileName, "r"); // If here, file exists so need for NULL check.
	doLoop1 = TRUE;
	while (doLoop1) {
		struct  individual *person = NULL;
		char phenotype[DATA_BUFFER_SIZE];

		// Genotype variables
		int gn = 0;
		int i = 0;
		BOOL linedone = FALSE;
		BOOL fatal = FALSE;
		char fmsg[msg_BUFFER_LENGTH];
		int f = 0; // Get first field

		if (feof(fin)) break;	

		// Read in the Person ID string.
		if (read_string(fin, buffer, DATA_BUFFER_SIZE)) f++;
		else break;

		// Named person so instantiate the person record for the input array.
		if (strlen(buffer) == 0) continue; // EOF hit?

		// Comment so skip current line and continue.
		if (buffer[0] == '#') {
			while(fgetc(fin) != '\n' && !feof(fin)) {} 
			continue;
		}

		// Selected the reserved "FID" ID code.
		if (strncmp(buffer, "FID", 3) == 0) {
			if(ops->verbose) error("FID is a reserved ID... please select a different family ID.");
			fclose(fin);
			if (ambiguous != NULL) free(ambiguous);
			return 0;
		}

		person = individual_create();
		if (person == NULL) return 0;
		individual_set_fid(person, buffer);

		// 6 or 7 obligatory fields.
		if (ops->pedSkipFid) {
			individual_set_iid(person, person->fid);
		} else {
			if (read_string(fin, buffer, DATA_BUFFER_SIZE)) f++;
			else break;
			individual_set_iid(person, buffer);
		}

		if (ops->pedSkipParents) {
			individual_set_pat(person, "0"); // default, junk values.
			individual_set_mat(person, "0");
		} else {
			// Paternal ID
			if (read_string(fin, buffer, DATA_BUFFER_SIZE)) f++;
			else break;
			individual_set_pat(person, buffer);

			// Maternal ID
			if (read_string(fin, buffer, DATA_BUFFER_SIZE)) f++;
			else break;
			individual_set_mat(person, buffer);
		}

		if (ops->pedSkipSex) 
			individual_set_sexcode(person, "0");
		else {
			if (read_string(fin, buffer, DATA_BUFFER_SIZE)) f++;
			else break;

			individual_set_sexcode(person, buffer);
		}

		if (ops->pedSkipPhenotype) {
			strcpy(phenotype, ops->missing_phenotype);
		} else {
			if (read_string(fin, phenotype, DATA_BUFFER_SIZE)) f++;
			else break;
		}

		// Are we using 0/1 coding.
		if (ops->coding01) {
			if (strcmp(phenotype, "1") == 0) strcpy(phenotype, "2");
			else if (strcmp(phenotype, "0") == 0) strcpy(phenotype, "1");
			else strcpy(phenotype, "0");
		}

		// Optional liability field, ignored by default, just read from the file.
		if (ops->liability) {
			if (read_string(fin, phenotype, DATA_BUFFER_SIZE)) f++;
			else break;
		}

		// Check sex
		if (strcmp(person->sexcode, "1") == 0) {
			person->sex = TRUE;
			nmale++;
		} else if (strcmp(person->sexcode, "2") == 0) { 
			person->sex = FALSE;
			nfemale++;
		} else {
			// Allocate the ambigious array as required.
			if (nAmbiguous == 0) {
				ambiguous = (struct individual**) malloc(sizeof(struct individual*) * INITIAL_ITEM_ARRAY_SIZE);
				if (ambiguous != NULL) {
					ambiguous_available = INITIAL_ITEM_ARRAY_SIZE;
				} else {
					fclose(fin);
					return 0;
				}
			} else if (nAmbiguous == ambiguous_available) {
				struct individual **pOld = NULL, **pNew = NULL;

				pOld = ambiguous;
				pNew = (struct individual**) realloc(ambiguous, (sizeof(struct individual*) * nAmbiguous) + (sizeof(struct individual*) * INITIAL_ITEM_ARRAY_SIZE));
				if (pNew != NULL) {
					ambiguous = pNew;
					ambiguous_available += INITIAL_ITEM_ARRAY_SIZE;
				} else {
					if(pOld != NULL) free(pOld);
					fclose(fin);
					return 0;
				}
			}

			// Add person to the ambigious array.
			ambiguous[nAmbiguous] = person;
			nAmbiguous++;

			if (!ops->ignore_missing_sex) person->missing = TRUE;
		}

		// A non-founder?
		person->founder = (strcmp(person->pat, "0") == 0 && strcmp(person->mat, "0") == 0) ? TRUE : FALSE;

		// Test for quantitative trait
		if (strcmp(phenotype, ops->missing_phenotype) == 0) {
			person->missing = TRUE;
		} else {
			person->phenotype = atof(phenotype);
			if ((int) person->phenotype != 0 && (int) person->phenotype != 1 &&
				(int) person->phenotype != 2) 
			{
				ops->qt = TRUE;
				ops->bt = FALSE;
			}
		}

		// necessary locus space.
		if (!ops->isSnpMajor) {
			// Memory allocation failure, bomb out.
			if (!individual_alloc_one(person, map->nLocus) || 
				!individual_alloc_two(person, map->nLocus)) {
					if (ambiguous != NULL) free(ambiguous);
					fclose(fin);
					return 0;
			}
		} 

		// Read genotypes now.
		while (!linedone) {
			BOOL doLoop2 = TRUE;
			struct locus *loc = NULL;
			char one[DATA_BUFFER_SIZE], two[DATA_BUFFER_SIZE];
			int nOne = 0, nTwo = 0;
			one[0] = two[0] = '\0';

			// Read in the 'one' string.
			while (doLoop2) {
				char ch = (char) fgetc(fin);
				if (ch==' ' || ch=='\t' || ch=='\n' || ch=='\r' || feof(fin)) 
				{
					if (ch=='\n' || ch=='\r' || feof(fin)) linedone = TRUE;
					if (nOne > 0) {
						gn++;
						break;
					}
					if (ch=='\n' || ch=='\r' || feof(fin)) break;
				} 
				else 
				{
					one[nOne] = ch;
					nOne++;
				}
			}
			one[nOne] = '\0';

			// Second allele
			if (!linedone) {
				// Read in the 'two' string.
				while (doLoop2) {
					char ch = (char) fgetc(fin);
					if (ch==' ' || ch=='\t' || ch=='\n' || ch=='\r' || feof(fin)) {
						if (ch=='\n' || ch=='\r' || feof(fin)) linedone = TRUE;
						if (nTwo > 0) {
							gn++;
							break;
						}
						if (ch=='\n' || ch=='\r' || feof(fin)) break;
					} else {
						two[nTwo] = ch;
						nTwo++;
					}
				}
			}
			two[nTwo] = '\0';

			if (linedone && nOne == 0 && nTwo == 0) break;

			loc = map->locus[i];

			// Add allele names to list, if needed
			// If allele is not missing.
			if (strcmp(one, ops->missing_genotype) != 0 && strcmp(one, ops->missing_genotype) != 0) {
				// ...and not already listed.
				if (strcmp(one, loc->allele1) != 0 && strcmp(one, loc->allele2) != 0) {
					if (strlen(loc->allele1) == 0) locus_set_allele1(loc, one);
					else if (strlen(loc->allele2) == 0) 
						locus_set_allele2(loc, one);
					else {
						// or show an error msg uf no empty slots.
						if (!fatal) {
							sprintf(fmsg, "this Locus %s has >2 alleles: \n individual %s %s has genotype [%s %s]\n but we've already seen [%s] and [%s]\n", 
								loc->name, person->fid, person->iid, one, two, loc->allele1, loc->allele2);
						}
						fatal = TRUE;
					} 
				} 
			} 

			// Repeat for the second allele, if different
			if (strcmp(one, two) != 0) {
				// If allele is not missing...
				if (strcmp(one, ops->missing_genotype) != 0) {
					// ...and not already listed.
					if (strcmp(two, loc->allele1) != 0 && strcmp(two, loc->allele2) != 0) {
						if (strlen(loc->allele1) == 0) locus_set_allele1(loc, two);
						else if (strlen(loc->allele2) == 0) locus_set_allele2(loc, two);
						else {
							// or show an error msg uf no empty slots.
							if (!fatal) {
								sprintf(fmsg, "Locus %s has >2 alleles: \n individual %s %s has genotype [%s %s]\n but we've already seen [%s] and [%s]\n", 
									loc->name, person->fid, person->iid, one, two, loc->allele1, loc->allele2);
							}
							fatal = TRUE; 
						} 
					} 
				}
			} 

			// Add specific genotypes.
			if (ops->isSnpMajor) {
				// 00 hom
				if (strcmp(one, loc->allele1) == 0 && strcmp(two, loc->allele1) == 0) {
					map->SNP[i]->one[map->SNP[i]->currentOneIndex++] = FALSE;
					map->SNP[i]->two[map->SNP[i]->currentTwoIndex++] = FALSE;
				}

				// 01 het
				else if (strcmp(one, ops->missing_genotype) != 0 && strcmp(two, ops->missing_genotype) != 0
					&& strcmp(one, two) != 0) 
				{
					map->SNP[i]->one[map->SNP[i]->currentOneIndex++] = FALSE;
					map->SNP[i]->two[map->SNP[i]->currentTwoIndex++] = TRUE;
				}

				// 11 hom
				else if (strcmp(one, loc->allele2) == 0 && strcmp(two, loc->allele2) == 0) {
					map->SNP[i]->one[map->SNP[i]->currentOneIndex++] = TRUE;
					map->SNP[i]->two[map->SNP[i]->currentTwoIndex++] = TRUE;
				}

				// 10 missing
				else if (strcmp(one, ops->missing_genotype) == 0 || strcmp(two, ops->missing_genotype) == 0) {
					map->SNP[i]->one[map->SNP[i]->currentOneIndex++] = TRUE;
					map->SNP[i]->two[map->SNP[i]->currentTwoIndex++] = FALSE;
				}
			} else {
				// 00 hom
				if (strcmp(one, loc->allele1) == 0 && strcmp(two, loc->allele1) == 0) {
					person->one[i] = FALSE;
					person->two[i] = FALSE;
				}

				// 01 het
				else if (strcmp(one, ops->missing_genotype) != 0 && strcmp(two, ops->missing_genotype) != 0
					&& strcmp(one, two) != 0) 
				{
					person->one[i] = FALSE;
					person->two[i] = TRUE;
				}

				// 11 hom
				else if (strcmp(one, loc->allele2) == 0 && strcmp(two, loc->allele2) == 0) {
					person->one[i] = TRUE;
					person->two[i] = TRUE;
				}

				// 10 missing
				else if (strcmp(one, ops->missing_genotype) == 0 || strcmp(two, ops->missing_genotype) == 0) {
					person->one[i] = TRUE;
					person->two[i] = FALSE;
				} 
			} 

			i++; // Advance to next locus.

			if (i > map->nLocus) {
				int ef = 1;

				if (!ops->pedSkipFid) ef++;
				if (!ops->pedSkipParents) ef += 2;
				if (!ops->pedSkipSex) ef++;
				if (!ops->pedSkipPhenotype) ef++;
				if (!ops->liability) ef++;

				// Close any hanging memory and exit, flagging an error.
				if (fin != NULL) fclose(fin);
				if (ambiguous != NULL) free(ambiguous);

				sprintf(fmsg, "Problem with line %i in [%s]\nExpecting %i + 2 * %i = %i columns, but found more.", 
					(c + 1), ops->szPedFileName, ef, map->nLocus, (ef +  (2 * map->nLocus)));
				if (ops->verbose) error(fmsg);

				return 0; 
			} 
		}
		
		// check size of the length somewhere
		if (gn != (2 * map->nLocus)) {
			int ef = 1;

			if (!ops->pedSkipFid) ef++;
			if (!ops->pedSkipParents) ef += 2;
			if (!ops->pedSkipSex) ef++;
			if (!ops->pedSkipPhenotype) ef++;
			if (!ops->liability) ef++;

			// Close any hanging memory and exit, flagging an error.
			if (fin != NULL) fclose(fin);
			if (ambiguous != NULL) free(ambiguous);

			sprintf(fmsg, "Problem with line %i in [%s]\nExpecting %i + 2 * %i = %i columns, but found %i.", 
				(c + 1), ops->szPedFileName, ef, map->nLocus, (ef +  (2 * map->nLocus)), (f + gn));

			fatal = TRUE;
		} 

		if (fatal) {
			if (ops->verbose) error(fmsg);
			return 0;
		}
		ped->sample[c++] = person;  // Increase person counter
	}
	if (fin != NULL) fclose(fin);

	if (ped->nSample > 0) ped->np = getPairwiseTests(ped->nSample);

	// If binary trait, now make 0 missing also
	// i.e. if we never saw other missing 0, 1, or 2.
	if (ops->bt) {
		int i;
		for(i = 0; i < ped->nSample; i++) {
			if (ped->sample[i] == NULL) continue;
			if (ped->sample[i]->phenotype == 0) ped->sample[i]->missing = TRUE;
			else ped->sample[i]->missing = FALSE;
		}
	} 
	if (ambiguous != NULL) free(ambiguous); 

	sprintf(msg, "%i individuals read from [ %s ]", ped->nSample, ops->szPedFileName);
	if(ops->verbose) print_log(msg);
	sprintf(msg, "%i males, %i females and 0 of unspecified sex", nmale, nfemale, nAmbiguous);
	if(ops->verbose) print_log(msg);
	
	// Echo number of missing phenotypes.
	for (j = 0; j < ped->nSample; j++) {
		if (ped->sample[j] == NULL) continue;
		else if (ped->sample[j]->missing) nm++;
	}
	sprintf(msg, "%i individuals with non-missing phenotypes", nm);
	if (ops->verbose) print_log(msg); 

	// By default, the things being analysed are binary traits.
	if (ops->bt) {
		int ncase = 0, ncontrol = 0, nmissing = 0, i;

		if (ops->coding01) {
			if (ops->verbose) print_log("Assuming a disease phenotype (0=unaff, 1=aff, other=miss)");
		} else {
			if(ops->verbose) print_log("Assuming a disease phenotype (1=unaff, 2=aff, 0=miss)");
			if (strcmp(ops->missing_phenotype, "0") != 0) {
				sprintf(msg, "Missing phenotype value is also %s", ops->missing_phenotype);
				if (ops->verbose) print_log(msg);
			}
		}

		for (i = 0; i < ped->nSample; i++) {
			if (ped->sample[i] == NULL) continue;
			if (ped->sample[i]->missing) nmissing++;
			else if (ped->sample[i]->phenotype == 1) ncontrol++;
			else if (ped->sample[i]->phenotype == 2) ncase++;
		}
		sprintf(msg, "%i case, %i controls and %i missing", ncase, ncontrol, nmissing);
		if (ops->verbose) print_log(msg);
	} else {
		sprintf(msg, "Missing phenotype value is %s", ops->missing_phenotype);
		if (ops->verbose) {
			print_log("Assuming a quantitative trait");
			print_log(msg);
		}
	} 

	return 1;
}

int pedfile_insert_sample(struct pedfile* ped, struct individual *person) {
	if (ped == NULL || person == NULL) {
		return 0;
	}

	// Alloc the required memory.
	if (ped->nSample == 0) {
		ped->sample = (struct individual**) malloc(sizeof(struct individual*) * INITIAL_ITEM_ARRAY_SIZE);
		if (ped->sample != NULL) {
			ped->sample_available = INITIAL_ITEM_ARRAY_SIZE;
		} else {
			return 0;
		}
	} else if (ped->nSample == ped->sample_available) {
		struct individual **pOld = NULL, **pNew = NULL;
		int new_size = 
			(sizeof(struct individual*) * ped->sample_available) + 
			(sizeof(struct individual*) * INITIAL_ITEM_ARRAY_SIZE);

		pOld = ped->sample;
		pNew = (struct individual**) realloc(ped->sample, new_size);
		if (pNew != NULL) {
			ped->sample = pNew;
			ped->sample_available += INITIAL_ITEM_ARRAY_SIZE;
		} else {
			ped->sample = pOld;
			return 0;
		}
	}

	// Add value to the array.
	ped->sample[ped->nSample] = person;
	ped->nSample++;

	return 1;
}

int pedfile_init(struct pedfile *ped) {
	if (ped != NULL)  {
		ped->sample = NULL;
		ped->nSample = 0;
		ped->sample_available = 0;
		ped->np = 0;
	}

	return 0;
}

int pedfile_free_attributes(struct pedfile *ped) {
	int i;

	if (ped == NULL) return 0;

	// Delete the person's array.
	if (ped->sample != NULL) {
		for (i = 0; i < ped->nSample; i++) {
			individual_free_attributes(ped->sample[i]);
			free(ped->sample[i]);
		}
		if (ped->sample != NULL) free(ped->sample);
		ped->sample = NULL;
		ped->nSample = 0;
		ped->np = 0;
		ped->sample_available = 0;
	}

	return 1;
}

struct pedfile* pedfile_create() {
	struct pedfile* p = NULL;

	p = (struct pedfile*) malloc(sizeof(struct pedfile));
	pedfile_init(p);

	return p;
}

int pedfile_alloc_sample(struct pedfile *ped, int size) {
	int i;	
	if (ped == NULL || size == 0) return 0;

	// Already allocated so multiple invocations would cause a memory leak.
	else if (ped->sample_available > 0) return 0;

	ped->sample = (struct individual**) malloc(sizeof(struct individual*) * size);
	if (ped->sample == NULL) return 0;
	else {
		for (i = 0; i < size; i++) ped->sample[i] = NULL;
		ped->nSample = size;
		ped->sample_available = size;
	}

	return 1;
}

// Create a record of the original order of the sample array 
// before permutation.
int pedfile_set_original_order(struct pedfile *ped) {
	int i;

	if (ped == NULL) return 0;
	else if (ped->nSample == 0) return 0;

	for (i = 0; i < ped->nSample; i++)
		ped->sample[i]->pperson = ped->sample[i];

	return 1;
}

int pedfile_aff_coding(struct pedfile *ped) {
	int i;

	if (ped == NULL) return 0;

	for (i = 0; i < ped->nSample; i++) {
		if (ped->sample[i]->phenotype == 2) 
			ped->sample[i]->aff = TRUE;
		else
			ped->sample[i]->aff = FALSE;
	}

	return 1;
}

int pedfile_copy_and_read(struct pedfile *p, struct mapfile *m, struct selected_options *ops, int rank) {
	BOOL jobDone = FALSE;
	char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH];

	if (m == NULL || ops == NULL || p == NULL) return 0;

	strcpy(remoteFilepath, ops->szPedFileName);
	set_datafile_path(localFilepath, ops->ldir, ops->szRunId, rank, SUFFIX_PED_FILE);
	if (copyFile(remoteFilepath, localFilepath)) {
		strcpy(ops->szPedFileName, localFilepath);
		if (pedfile_read(p, m, ops)) 
			if (deleteFile(localFilepath)) jobDone = TRUE;
	}

	if (jobDone) return 1;
	else return 0;
}